# 存储模块
import pandas as pd
import json
import os
import sys
from datetime import datetime

# 获取当前路径（兼容PyInstaller打包后的情况）
def get_current_path():
    if hasattr(sys, '_MEIPASS'):
        return os.path.dirname(os.path.abspath(sys.executable))
    return os.path.dirname(os.path.abspath(__file__))

# 导入配置模块
try:
    from config import OUTPUT_FORMAT, OUTPUT_DIRECTORY, MAIN_FORCE_FILENAME
except ImportError:
    # 如果无法导入配置，使用默认值
    print("警告：无法导入config模块，使用默认配置")
    OUTPUT_FORMAT = 'excel'
    OUTPUT_DIRECTORY = 'output'
    MAIN_FORCE_FILENAME = 'main_force_data'

class DataStorage:
    """数据存储类"""
    
    def __init__(self):
        self.data = []
        
    def add_data(self, news_item):
        """添加数据到内存"""
        if self._validate_data(news_item):
            # 添加时间戳
            news_item['crawl_time'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            self.data.append(news_item)
            return True
        return False
    
    def _validate_data(self, data_item):
        """验证数据完整性"""
        # 主力排名数据的基本验证字段
        required_fields = ['stock_code', 'stock_name']
        for field in required_fields:
            if field not in data_item or not data_item[field]:
                return False
        return True
    
    def save_to_excel(self, filename=None):
        """保存数据到Excel文件"""
        if not self.data:
            print("警告：没有数据需要保存")
            return False
            
        # 获取当前路径并构建输出目录路径
        current_dir = get_current_path()
        output_dir = os.path.join(current_dir, OUTPUT_DIRECTORY)
        
        # 确保输出目录存在
        try:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                print(f"已创建输出目录: {output_dir}")
        except Exception as e:
            print(f"创建输出目录失败: {e}")
            return False
            
        # 添加日期时间标识到文件名
        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        if filename:
            base_name, ext = os.path.splitext(filename)
            filename = f"{base_name}_{current_time}{ext}"
        else:
            filename = f"{MAIN_FORCE_FILENAME}_{current_time}.xlsx"
        
        filepath = os.path.join(output_dir, filename)
        
        try:
            df = pd.DataFrame(self.data)
            
            # 表头映射：英文字段名转中文
            header_mapping = {
                'stock_code': '股票代码',
                'stock_name': '股票名称',
                'current_price': '当前价格',
                'price_change_rate': '涨跌幅',
                'turnover_rate': '换手率',
                'turnover_amount': '成交额',
                'net_inflow': '净流入',
                'main_force_net_inflow': '主力净流入',
                'main_force_ratio': '主力占比',
                'rank': '排名',
                'crawl_time': '爬取时间',
                'source_url': '数据来源',
                'symbol': '股票代码',
                'name': '股票名称',
                'trade': '当前价格',
                'changeratio': '涨跌幅',
                'turnover': '换手率',
                'amount': '成交额',
                'netamount': '净流入',
                'r0_net': '主力净流入',
                'r0_ratio': '主力占比'
            }
            
            # 重命名列名
            df.rename(columns=header_mapping, inplace=True)
            
            # 保存到Excel
            df.to_excel(filepath, index=False, engine='openpyxl')
            print(f"数据已保存到Excel文件：{filepath}")
            return True
        except Exception as e:
            print(f"保存Excel文件失败：{e}")
            return False
    
    def save_to_json(self, filename=None):
        """保存数据到JSON文件"""
        if not self.data:
            print("警告：没有数据需要保存")
            return False
            
        # 获取当前路径并构建输出目录路径
        current_dir = get_current_path()
        output_dir = os.path.join(current_dir, OUTPUT_DIRECTORY)
        
        # 确保输出目录存在
        try:
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                print(f"已创建输出目录: {output_dir}")
        except Exception as e:
            print(f"创建输出目录失败: {e}")
            return False
            
        # 添加日期时间标识到文件名
        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        if filename:
            base_name, ext = os.path.splitext(filename)
            filename = f"{base_name}_{current_time}{ext}"
        else:
            filename = f"{MAIN_FORCE_FILENAME}_{current_time}.json"
        
        filepath = os.path.join(output_dir, filename)
        
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(self.data, f, ensure_ascii=False, indent=2)
            print(f"数据已保存到JSON文件：{filepath}")
            return True
        except Exception as e:
            print(f"保存JSON文件失败：{e}")
            return False
    
    def save_data(self, filename=None):
        """根据配置保存数据"""
        if OUTPUT_FORMAT == 'excel':
            return self.save_to_excel(filename)
        elif OUTPUT_FORMAT == 'json':
            return self.save_to_json(filename)
        else:
            print(f"不支持的输出格式：{OUTPUT_FORMAT}")
            return False
    
    def get_data_count(self):
        """获取数据数量"""
        return len(self.data)
    
    def clear_data(self):
        """清空内存中的数据"""
        self.data.clear()